/*
 * Copyright 2023-2025, Haiku, Inc. All rights reserved.
 * Distributed under the terms of the MIT License.
 *
 * 龙芯LoongArch64架构内存管理单元(MMU)实现
 * 本文件实现LoongArch64架构的MMU相关功能，包括TLB管理和地址转换
 */

#include <KernelExport.h>

#include <arch/cpu.h>
#include <arch/mmu.h>
#include <boot/kernel_args.h>
#include <kernel.h>
#include <vm/vm.h>
#include <vm/VMAddressSpace.h>

#include "LoongArchVMTranslationMap.h"

// 定义调试输出宏
#define TRACE_MMU
#ifdef TRACE_MMU
#define TRACE(x...) dprintf("arch_mmu: " x)
#else
#define TRACE(x...) ;
#endif

// 定义LoongArch64 TLB相关CSR寄存器
#define CSR_TLBIDX    0x10    // TLB索引
#define CSR_TLBEHI    0x11    // TLB入口高位
#define CSR_TLBELO0   0x12    // TLB入口低位0
#define CSR_TLBELO1   0x13    // TLB入口低位1
#define CSR_ASID      0x18    // 地址空间标识符
#define CSR_PGDL      0x19    // 页表基址低位
#define CSR_PGDH      0x1a    // 页表基址高位
#define CSR_PGD       0x1b    // 页表基址

// 定义TLB操作指令
#define TLBSRCH     0x0       // 搜索TLB
#define TLBRD       0x1       // 读取TLB
#define TLBWR       0x2       // 写入TLB
#define TLBFILL     0x3       // 填充TLB
#define TLBCLR      0x4       // 清除TLB
#define TLBFLUSH    0x5       // 刷新TLB
#define INVTLB      0x6       // 无效化TLB

// 定义TLBIDX寄存器位
#define TLBIDX_PS_SHIFT  24    // 页大小偏移
#define TLBIDX_PS_MASK   0x3f  // 页大小掩码
#define TLBIDX_INDEX_SHIFT 0   // 索引偏移
#define TLBIDX_INDEX_MASK 0xfff // 索引掩码

// 定义INVTLB操作码
#define INVTLB_ALL          0    // 无效化所有TLB
#define INVTLB_ASID         1    // 无效化指定ASID的TLB
#define INVTLB_VA           2    // 无效化指定VA的TLB
#define INVTLB_ASID_VA      3    // 无效化指定ASID和VA的TLB
#define INVTLB_ASID_VA_PAGE 4    // 无效化指定ASID和VA页的TLB

// 定义页大小
#define PS_4K   12   // 4KB页
#define PS_8K   13   // 8KB页
#define PS_16K  14   // 16KB页
#define PS_32K  15   // 32KB页
#define PS_64K  16   // 64KB页
#define PS_128K 17   // 128KB页
#define PS_256K 18   // 256KB页
#define PS_512K 19   // 512KB页
#define PS_1M   20   // 1MB页
#define PS_2M   21   // 2MB页
#define PS_4M   22   // 4MB页
#define PS_8M   23   // 8MB页
#define PS_16M  24   // 16MB页
#define PS_32M  25   // 32MB页
#define PS_64M  26   // 64MB页
#define PS_128M 27   // 128MB页
#define PS_256M 28   // 256MB页
#define PS_512M 29   // 512MB页
#define PS_1G   30   // 1GB页
#define PS_2G   31   // 2GB页
#define PS_4G   32   // 4GB页

// 全局变量
static uint32 sTlbEntries = 0;  // TLB条目数量
static uint32 sTlbPageSizes = 0; // 支持的页大小

// 初始化MMU
status_t
arch_mmu_init(kernel_args *args)
{
    TRACE("Initializing MMU\n");

    // 检测TLB条目数量
    // 通过读取CPUCFG寄存器获取TLB信息
    uint32 cpucfg2 = read_cpucfg(2);
    sTlbEntries = ((cpucfg2 >> 16) & 0x3ff) + 1;
    
    // 检测支持的页大小
    uint32 cpucfg3 = read_cpucfg(3);
    sTlbPageSizes = (cpucfg3 >> 8) & 0xff;

    TRACE("TLB entries: %u, supported page sizes: 0x%x\n", sTlbEntries, sTlbPageSizes);

    return B_OK;
}

// 初始化每个CPU的MMU
status_t
arch_mmu_init_percpu(kernel_args *args, int cpu)
{
    TRACE("Initializing MMU for CPU %d\n", cpu);

    // 设置页表基址
    uint64 pgd = args->arch_args.phys_pgd;
    asm volatile("csrwr %0, %1" : : "r"(pgd), "i"(CSR_PGD) : "memory");

    // 清除TLB
    asm volatile("invtlb %0, $zero, $zero" : : "i"(INVTLB_ALL) : "memory");

    return B_OK;
}

// 无效化TLB范围
void
arch_cpu_invalidate_TLB_range(addr_t start, addr_t end)
{
    // 确保地址对齐到页边界
    start = ROUNDDOWN(start, B_PAGE_SIZE);
    end = ROUNDUP(end, B_PAGE_SIZE);

    // 遍历范围内的每个页面，无效化TLB
    for (addr_t addr = start; addr < end; addr += B_PAGE_SIZE) {
        asm volatile("invtlb %0, $zero, %1" : : "i"(INVTLB_VA), "r"(addr) : "memory");
    }
}

// 无效化TLB列表
void
arch_cpu_invalidate_TLB_list(addr_t pages[], int num_pages)
{
    for (int i = 0; i < num_pages; i++) {
        asm volatile("invtlb %0, $zero, %1" : : "i"(INVTLB_VA), "r"(pages[i]) : "memory");
    }
}

// 全局无效化TLB
void
arch_cpu_global_TLB_invalidate(void)
{
    asm volatile("invtlb %0, $zero, $zero" : : "i"(INVTLB_ALL) : "memory");
}

// 用户空间无效化TLB
void
arch_cpu_user_TLB_invalidate(void)
{
    // 获取当前ASID
    uint32 asid;
    asm volatile("csrrd %0, %1" : "=r"(asid) : "i"(CSR_ASID));
    
    // 无效化指定ASID的所有TLB条目
    asm volatile("invtlb %0, %1, $zero" : : "i"(INVTLB_ASID), "r"(asid) : "memory");
}

// 设置页表基址
void
arch_mmu_set_page_directory(VMAddressSpace *aspace, phys_addr_t pgdPhysical)
{
    if (aspace == VMAddressSpace::Kernel()) {
        // 设置内核页表基址
        asm volatile("csrwr %0, %1" : : "r"(pgdPhysical), "i"(CSR_PGD) : "memory");
    } else {
        // 用户空间页表基址在上下文切换时设置
        // 这里只保存页表基址，不立即设置
        aspace->TranslationMap()->SetPageTable(pgdPhysical);
    }
}

// 获取页表基址
phys_addr_t
arch_mmu_get_page_directory(VMAddressSpace *aspace)
{
    if (aspace == VMAddressSpace::Kernel()) {
        // 获取内核页表基址
        uint64 pgd;
        asm volatile("csrrd %0, %1" : "=r"(pgd) : "i"(CSR_PGD));
        return pgd;
    } else {
        // 获取用户空间页表基址
        return aspace->TranslationMap()->PageTable();
    }
}

// 手动填充TLB条目
status_t
arch_mmu_fill_tlb(addr_t va, phys_addr_t pa, uint32 attributes, uint32 memoryType)
{
    // 构建TLB条目
    uint64 tlbelo0 = (pa & kPtePhysAddrMask);
    
    // 设置访问权限
    if ((attributes & B_READ_AREA) != 0)
        tlbelo0 |= kPteRead;
    if ((attributes & B_WRITE_AREA) != 0 || (attributes & B_KERNEL_WRITE_AREA) != 0)
        tlbelo0 |= kPteWrite;
    if ((attributes & B_EXECUTE_AREA) != 0)
        tlbelo0 |= kPteExec;

    // 设置用户/内核访问权限
    if ((attributes & B_USER_PROTECTION) != 0)
        tlbelo0 |= kPteUser;

    // 设置缓存属性
    if (memoryType == B_MTR_UC)
        tlbelo0 |= kPteNoCache;

    // 设置全局标志（内核页面通常是全局的）
    if ((attributes & B_KERNEL_AREA) != 0)
        tlbelo0 |= kPteGlobal;

    // 设置有效位和访问/脏位
    tlbelo0 |= kPteValid | kPteAccessed | kPteDirty;

    // 设置TLB条目
    asm volatile(
        "csrwr %0, %3\n"  // 设置虚拟地址
        "csrwr %1, %4\n"  // 设置物理地址和属性
        "csrwr %1, %5\n"  // 两个页共用同一个设置
        "tlbfill\n"       // 填充TLB
        : 
        : "r"(va), "r"(tlbelo0), "i"(PS_4K), 
          "i"(CSR_TLBEHI), "i"(CSR_TLBELO0), "i"(CSR_TLBELO1)
        : "memory"
    );

    return B_OK;
}